package dev.alm.mcpdemo.config;

import dev.alm.mcpdemo.service.ChatService;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.mcp.McpToolProvider;
import dev.langchain4j.mcp.client.DefaultMcpClient;
import dev.langchain4j.mcp.client.transport.stdio.StdioMcpTransport;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;
import java.util.Map;

@Configuration
public class LLMConfig {

    @Value("${llm.api.key}")
    private String apiKey;

    @Value("${llm.base.url}")
    private String baseUrl;

    @Value("${llm.model.name}")
    private String modelName;

    @Bean
    public OpenAiChatModel openAiChatModel() {
        return OpenAiChatModel.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .modelName(modelName)
                .build();
    }

    @Bean
    public OpenAiStreamingChatModel openAiStreamingChatModel() {
        return OpenAiStreamingChatModel.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .modelName(modelName)
                .build();
    }

    @Bean
    public ChatService chatService(OpenAiChatModel openAiChatModel, OpenAiStreamingChatModel openAiStreamingChatModel) {
        // 构建McpTransport协议
        StdioMcpTransport transport = new StdioMcpTransport.Builder()
                .command(List.of("cmd", "/c", "npx", "-y", "@baidumap/mcp-server-baidu-map"))
//                .environment(Map.of("BAIDU_MAP_API_KEY", System.getenv("BAIDU_MAP_API_KEY")))
                .environment(Map.of("BAIDU_MAP_API_KEY", "abc123"))
                .build();

        // 构建McpClient客户端
        DefaultMcpClient mcpClient = new DefaultMcpClient.Builder().transport(transport).build();

        // 创建工具集
        // 和原生Function Calling类似
        McpToolProvider mcpToolProvider = McpToolProvider.builder().mcpClients(mcpClient).build();
        ChatService service = AiServices.builder(ChatService.class)
                .chatModel(openAiChatModel)
                .streamingChatModel(openAiStreamingChatModel)
                .toolProvider(mcpToolProvider)
                .build();
        return service;
    }

}
